/*
* Copyright (c) 2014 Google, Inc.
*
* This software is provided 'as-is', without any express or implied
* warranty.  In no event will the authors be held liable for any damages
* arising from the use of this software.
* Permission is granted to anyone to use this software for any purpose,
* including commercial applications, and to alter it and redistribute it
* freely, subject to the following restrictions:
* 1. The origin of this software must not be misrepresented; you must not
* claim that you wrote the original software. If you use this software
* in a product, an acknowledgment in the product documentation would be
* appreciated but is not required.
* 2. Altered source versions must be plainly marked as such, and must not be
* misrepresented as being the original software.
* 3. This notice may not be removed or altered from any source distribution.
*/
#ifndef B2_SLAB_ALLOCATOR_H
#define B2_SLAB_ALLOCATOR_H

#include <stddef.h>
#include <stdint.h>
#include <new>
#include "b2_intrusive_list.h"
#include "b2_free_list.h"
#include "b2_settings.h"
#include "b2_tracked_block.h"

/// Freelist based allocator for fixed sized items from slabs (memory
/// preallocated from the heap).
/// T should be a class which has a default constructor and implements the
/// member function "b2IntrusiveList* GetListNode()".
/// All objects in a slab are constructed when a slab is created and destructed
/// when a slab is freed.
template<typename T>
class b2SlabAllocator
{
private:
  // Information about a slab.
  class Slab
  {
  public:
    /// Initialize a slab with the number of items it contains.
    Slab(uint32 numberOfItems) :
      m_numberOfItems(numberOfItems)
    {
      B2_NOT_USED(m_padding);
      // This assumes that this class is packed on at least a 4-byte
      // boundary with no padding.  Verify the assumption.
      b2Assert(sizeof(*this) == b2_mallocAlignment);
    }

    /// Empty destructor.
    ~Slab() { }

    /// Get the number of items in this slab.
    uint32 GetNumberOfItems() const { return m_numberOfItems; }

    /// Get a pointer to the first item in the slab.
    T* GetFirstItem() const
    {
      return (T*)((uint8*)(this + 1));
    }

    /// Get a pointer to the end of the slab.
    /// NOTE: This is a pointer after the last byte of the slab not the
    /// last item in the slab.
    T* GetItemEnd() const { return GetFirstItem() + GetNumberOfItems(); }

  private:
    /// Number of items in the slab.
    uint32 m_numberOfItems;
    /// Padding to align the first item in the slab to b2_mallocAlignment.
    uint8 m_padding[b2_mallocAlignment - sizeof(uint32)];
  };

public:
  /// Initialize the allocator to allocate itemsPerSlab of type T for each
  /// slab that is allocated.
  b2SlabAllocator(const uint32 itemsPerSlab) :
    m_itemsPerSlab(itemsPerSlab)
  {
  }

  /// Free all allocated slabs.
  ~b2SlabAllocator()
  {
    FreeAllSlabs();
  }

  /// Set size of the next allocated slab using the number of items per
  /// slab.  Setting this value to zero disables further slab allocation.
  void SetItemsPerSlab(uint32 itemsPerSlab)
  {
    m_itemsPerSlab = itemsPerSlab;
  }

  // Get the size of the next allocated slab.
  uint32 GetItemsPerSlab() const
  {
    return m_itemsPerSlab;
  }

  /// Allocate a item from the slab.
  T* Allocate()
  {
    // Allocate a slab if needed here.
    if (m_freeList.GetFreeList()->GetFreeList().IsEmpty() &&
      !AllocateSlab())
      return NULL;
    return m_freeList.Allocate();
  }

  /// Free an item from the slab.
  void Free(T *object)
  {
    m_freeList.Free(object);
  }

  /// Allocate a slab, construct instances of T and add them to the free
  /// pool.
  bool AllocateSlab()
  {
    if (!m_itemsPerSlab) return false;
    const uint32 slabSize = sizeof(Slab) + (sizeof(T) * m_itemsPerSlab);
    void* const memory = m_slabs.Allocate(slabSize);
    if (!memory) return false;

    Slab* const slab = new (BlockGetSlab(memory)) Slab(m_itemsPerSlab);
    T* item = slab->GetFirstItem();
    for (uint32 i = 0; i < m_itemsPerSlab; ++i, ++item)
    {
      m_freeList.AddToFreeList(new (item) T);
    }
    return true;
  }

  /// Free all slabs.
  void FreeAllSlabs()
  {
    const b2TypedIntrusiveListNode<b2TrackedBlock>& slabList =
      m_slabs.GetList();
    while (!slabList.IsEmpty())
    {
      FreeSlab(BlockGetSlab(slabList.GetNext()->GetMemory()));
    }
  }

  /// Free all empty slabs.
  /// This method is slow - O(M^N) - since this class doesn't track
  /// the association between each item and slab.
  void FreeEmptySlabs()
  {
    const b2IntrusiveListNode& freeItemList =
      m_freeList.GetFreeList()->GetFreeList();
    const b2IntrusiveListNode* freeItemListTerminator =
      freeItemList.GetTerminator();
    const b2TypedIntrusiveListNode<b2TrackedBlock>& slabList =
      m_slabs.GetList();
    const b2TypedIntrusiveListNode<b2TrackedBlock>* slabListTerminator =
      slabList.GetTerminator();
    b2TrackedBlock* block = slabList.GetNext();
    while (block != slabListTerminator)
    {
      // Get the Slab from the memory associated with the block.
      Slab* const slab = BlockGetSlab(block->GetMemory());
      block = block->GetNext();

      // Determine the range of memory the Slab owns.
      const uint8* const slabItemStart = (uint8*)slab->GetFirstItem();
      const uint8* const slabItemEnd = (uint8*)slab->GetItemEnd();

      // Count all free items that are owned by the current slab.
      uint8 freeItems = 0;
      bool empty = false;
      for (b2IntrusiveListNode* itemNode = freeItemList.GetNext();
         itemNode != freeItemListTerminator;
         itemNode = itemNode->GetNext())
      {
        const uint8* itemNodeAddress = (uint8*)itemNode;
        if (itemNodeAddress >= slabItemStart &&
          itemNodeAddress <= slabItemEnd)
        {
          ++freeItems;
          if (slab->GetNumberOfItems() == freeItems)
          {
            empty = true;
            break;
          }
        }
      }
      // If a slab is empty, free it.
      if (empty)
      {
        FreeSlab(slab);
      }
    }
  }

  /// Get the item allocator freelist.
  const b2TypedFreeList<T>& GetFreeList() const
  {
    return m_freeList;
  }

private:
  /// Destroy all objects in a slab and free the slab.
  void FreeSlab(Slab * const slab)
  {
    b2Assert(slab);
    const uint32 numberOfItems = slab->GetNumberOfItems();
    T* item = slab->GetFirstItem();
    for (uint32 i = 0; i < numberOfItems; ++i, ++item)
    {
      item->~T();
    }
    slab->~Slab();
    m_slabs.Free(slab);
  }

  /// Get a pointer to a Slab from a block of memory in m_slabs.
  Slab* BlockGetSlab(void *memory)
  {
    return (Slab*)memory;
  }

  /// Get a pointer to the first item in the array of items referenced by a
  /// Slab.
  T* SlabGetFirstItem(Slab* slab)
  {
    return (T*)(slab + 1);
  }

private:
  /// Contains a list of b2TrackedBlock instances where each b2TrackedBlock's
  /// associated user memory contains a Slab followed by instances of T.
  b2TrackedBlockAllocator m_slabs;
  /// Number of items to allocate in the next allocated slab.
  uint32 m_itemsPerSlab;
  /// Freelist which contains instances of T.
  b2TypedFreeList<T> m_freeList;
};

#endif  // B2_SLAB_ALLOCATOR_H
